#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
基金数据下载脚本
功能：
1. 调用fund_name_em接口下载基金清单并存储到fund_info表
2. 调用fund_individual_basic_info_xq接口下载基金详细信息并存储到fund_info_detail表
http://localhost:3000/fund-data-downloader 实现基金列表下载和基金详细信息下载
"""
# type: ignore

import os
import sys
import django  # type: ignore
import akshare as ak  # type: ignore
import pandas as pd  # type: ignore
from datetime import datetime, date
import logging
from typing import Optional, Dict, Any
import time
from enum import Enum
import threading

# 设置Django环境
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, current_dir)
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'wealth_backend.settings')
django.setup()

from data_service.models import FundInfo, FundInfoDetail
from django.db import transaction  # type: ignore
from django.utils import timezone  # type: ignore

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('log/fund_data_download.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

class DownloadStatus(Enum):
    """下载状态枚举"""
    WAITING = "waiting"  # 等待下载
    DOWNLOADING = "downloading"  # 正在下载
    STOPPED = "stopped"  # 停止下载
    COMPLETED = "completed"  # 下载完成
    FAILED = "failed"  # 下载失败

class FundDataDownloader:
    """基金数据下载器"""
    
    def __init__(self):
        self.session_date = datetime.now().date()
        self._status = DownloadStatus.WAITING
        self._stop_flag = threading.Event()
        self._lock = threading.Lock()
        self._current_task = None
        self._progress_info = {
            'current_step': '',
            'current_index': 0,
            'total_count': 0,
            'success_count': 0,
            'failed_count': 0
        }
        logger.info(f"初始化基金数据下载器，会话日期：{self.session_date}")
    
    @property
    def status(self) -> DownloadStatus:
        """获取当前下载状态"""
        with self._lock:
            return self._status
    
    def set_status(self, status: DownloadStatus) -> None:
        """设置下载状态"""
        with self._lock:
            self._status = status
            logger.info(f"下载状态更新为：{status.value}")
    
    def stop_download(self) -> bool:
        """停止下载"""
        try:
            with self._lock:
                if self._status == DownloadStatus.DOWNLOADING:
                    self._stop_flag.set()
                    self._status = DownloadStatus.STOPPED
                    logger.info("下载停止请求已发送")
                    return True
                else:
                    logger.warning(f"当前状态为{self._status.value}，无法停止下载")
                    return False
        except Exception as e:
            logger.error(f"停止下载失败：{str(e)}")
            return False
    
    def reset(self) -> None:
        """重置下载器状态"""
        with self._lock:
            self._stop_flag.clear()
            self._status = DownloadStatus.WAITING
            self._current_task = None
            self._progress_info = {
                'current_step': '',
                'current_index': 0,
                'total_count': 0,
                'success_count': 0,
                'failed_count': 0
            }
        logger.info("下载器状态已重置")
    
    def is_stopped(self) -> bool:
        """检查是否已停止"""
        return self._stop_flag.is_set()
    
    def reset_stop_flag(self) -> None:
        """重置停止标志"""
        self._stop_flag.clear()
    
    def get_progress_info(self) -> Dict[str, Any]:
        """获取当前进度信息"""
        with self._lock:
            return {
                'status': self._status.value,
                'current_task': self._current_task,
                'progress': self._progress_info.copy()
            }
    
    def _update_progress(self, current_step: str, current_index: int = 0, total_count: int = 0, 
                        success_count: int = 0, failed_count: int = 0) -> None:
        """更新进度信息"""
        with self._lock:
            self._progress_info.update({
                'current_step': current_step,
                'current_index': current_index,
                'total_count': total_count,
                'success_count': success_count,
                'failed_count': failed_count
            })
    
    def download_fund_list(self) -> bool:
        """
        下载基金清单列表
        使用akshare的fund_name_em接口
        """
        try:
            # 设置状态和任务信息
            self.set_status(DownloadStatus.DOWNLOADING)
            self._current_task = "基金清单下载"
            self.reset_stop_flag()
            self._update_progress("开始下载基金清单")
            
            logger.info("开始下载基金清单列表...")
            
            # 检查是否被停止
            if self.is_stopped():
                logger.info("下载已被停止")
                self.set_status(DownloadStatus.STOPPED)
                return False
            
            # 调用akshare接口获取基金列表
            self._update_progress("正在获取基金列表数据")
            fund_list_df = ak.fund_name_em()
            
            if fund_list_df.empty:
                logger.error("获取的基金列表为空")
                self.set_status(DownloadStatus.FAILED)
                return False
            
            logger.info(f"获取到 {len(fund_list_df)} 只基金")
            
            # 检查是否被停止
            if self.is_stopped():
                logger.info("下载已被停止")
                self.set_status(DownloadStatus.STOPPED)
                return False
            
            # 数据清洗和映射
            self._update_progress("正在处理基金列表数据", total_count=len(fund_list_df))
            fund_list_df = fund_list_df.rename(columns={
                '基金代码': 'fund_code',
                '基金简称': 'fund_name', 
                '基金类型': 'fund_type',
                '拼音缩写': 'pinyin_abbr',
                '拼音全称': 'pinyin_full'
            })
            
            # 检查是否被停止
            if self.is_stopped():
                logger.info("下载已被停止")
                self.set_status(DownloadStatus.STOPPED)
                return False
            
            # 批量插入数据库
            self._update_progress("正在保存基金列表到数据库")
            with transaction.atomic():
                # 清空现有数据
                FundInfo.objects.all().delete()
                logger.info("已清空fund_info表中的现有数据")
                
                # 批量创建新记录
                fund_objects = []
                for _, row in fund_list_df.iterrows():
                    # 检查是否被停止
                    if self.is_stopped():
                        logger.info("下载已被停止")
                        self.set_status(DownloadStatus.STOPPED)
                        return False
                    
                    fund_obj = FundInfo(
                        fund_code=str(row.get('fund_code', '')),
                        fund_name=str(row.get('fund_name', '')),
                        fund_type=str(row.get('fund_type', '')),
                        fund_company='',  # 基金清单接口不包含公司信息
                        fund_manager='',  # 基金清单接口不包含经理信息
                        status='active',
                        data_source='akshare'
                    )
                    fund_objects.append(fund_obj)
                
                # 批量插入
                FundInfo.objects.bulk_create(fund_objects, batch_size=1000)
                logger.info(f"成功插入 {len(fund_objects)} 条基金清单记录")
            
            self._update_progress("基金清单下载完成", success_count=len(fund_objects))
            self.set_status(DownloadStatus.COMPLETED)
            return True
            
        except Exception as e:
            logger.error(f"下载基金清单失败：{str(e)}")
            return False
    
    def download_fund_details(self, limit: Optional[int] = None) -> bool:
        """
        下载基金详细信息
        使用akshare的fund_individual_basic_info_xq接口
        
        Args:
            limit: 限制下载数量，用于测试
        """
        try:
            # 设置状态和任务信息
            self.set_status(DownloadStatus.DOWNLOADING)
            self._current_task = "基金详细信息下载"
            self.reset_stop_flag()
            
            logger.info("开始下载基金详细信息...")
            
            # 获取所有基金代码
            self._update_progress("正在读取基金代码")
            logger.info("开始从数据库fund_info表中读取基金代码...")
            fund_codes = FundInfo.objects.values_list('fund_code', flat=True)
            logger.info(f"成功从fund_info表中读取到 {len(fund_codes)} 个基金代码")
            
            # 检查是否被停止
            if self.is_stopped():
                logger.info("下载已被停止")
                self.set_status(DownloadStatus.STOPPED)
                return False
            
            # 记录前10个基金代码作为示例
            if fund_codes:
                sample_codes = list(fund_codes[:10])
                logger.info(f"基金代码示例（前10个）: {sample_codes}")
            else:
                logger.warning("fund_info表中没有找到任何基金代码数据")
                self.set_status(DownloadStatus.FAILED)
                return False
            
            if limit:
                fund_codes = fund_codes[:limit]
                logger.info(f"限制下载数量为 {limit} 只基金")
            
            total_count = len(fund_codes)
            logger.info(f"需要下载 {total_count} 只基金的详细信息")
            
            # 检查是否被停止
            if self.is_stopped():
                logger.info("下载已被停止")
                self.set_status(DownloadStatus.STOPPED)
                return False
            
            # 清空现有详细信息数据
            self._update_progress("正在清理现有数据", total_count=total_count)
            with transaction.atomic():
                FundInfoDetail.objects.all().delete()
                logger.info("已清空fund_info_detail表中的现有数据")
            
            success_count = 0
            failed_count = 0
            
            for i, fund_code in enumerate(fund_codes, 1):
                try:
                    # 检查是否被停止
                    if self.is_stopped():
                        logger.info("下载已被停止")
                        self.set_status(DownloadStatus.STOPPED)
                        return False
                    
                    # 更新进度
                    self._update_progress(
                        f"正在下载第 {i}/{total_count} 只基金详细信息",
                        current_index=i,
                        total_count=total_count,
                        success_count=success_count,
                        failed_count=failed_count
                    )
                    
                    logger.info(f"正在下载第 {i}/{total_count} 只基金：{fund_code}")
                    
                    # 调用akshare接口获取基金详细信息，添加重试机制
                    fund_detail_df = None
                    max_retries = 3
                    
                    for retry in range(max_retries):
                        try:
                            # 检查是否被停止
                            if self.is_stopped():
                                logger.info("下载已被停止")
                                self.set_status(DownloadStatus.STOPPED)
                                return False
                            
                            fund_detail_df = ak.fund_individual_basic_info_xq(symbol=fund_code)
                            break
                        except Exception as retry_error:
                            logger.warning(f"基金 {fund_code} 第 {retry + 1} 次尝试失败：{str(retry_error)}")
                            if retry < max_retries - 1:
                                time.sleep(1)  # 等待1秒后重试
                            else:
                                raise retry_error
                    
                    if fund_detail_df is None or fund_detail_df.empty:
                        logger.warning(f"基金 {fund_code} 详细信息为空")
                        failed_count += 1
                        continue
                    
                    # 解析数据
                    detail_dict = {}
                    for _, row in fund_detail_df.iterrows():
                        detail_dict[row['item']] = row['value']
                    
                    # 创建详细信息记录
                    fund_detail = FundInfoDetail(
                        fund_code=fund_code,
                        fund_name=detail_dict.get('基金名称', ''),
                        fund_full_name=detail_dict.get('基金全称', ''),
                        establishment_date=self._parse_date(detail_dict.get('成立时间', '')),
                        latest_scale=detail_dict.get('最新规模', ''),
                        fund_company=detail_dict.get('基金公司', ''),
                        fund_manager=detail_dict.get('基金经理', ''),
                        custodian_bank=detail_dict.get('托管银行', ''),
                        fund_type=detail_dict.get('基金类型', ''),
                        rating_agency=detail_dict.get('评级机构', ''),
                        fund_rating=detail_dict.get('基金评级', ''),
                        investment_strategy=detail_dict.get('投资策略', ''),
                        investment_objective=detail_dict.get('投资目标', ''),
                        performance_benchmark=detail_dict.get('业绩比较基准', ''),
                        data_source='akshare'
                    )
                    
                    fund_detail.save()
                    success_count += 1
                    
                    # 避免请求过于频繁，增加延迟时间
                    time.sleep(0.5)  # 增加到0.5秒间隔
                    
                except Exception as e:
                    logger.error(f"下载基金 {fund_code} 详细信息失败：{str(e)}")
                    failed_count += 1
                    continue
            
            logger.info(f"基金详细信息下载完成，成功：{success_count}，失败：{failed_count}")
            
            # 更新最终进度
            self._update_progress(
                "基金详细信息下载完成",
                current_index=total_count,
                total_count=total_count,
                success_count=success_count,
                failed_count=failed_count
            )
            
            # 如果失败数量过多，认为下载失败
            if failed_count > success_count:
                logger.error(f"下载失败数量({failed_count})超过成功数量({success_count})，认为下载失败")
                self.set_status(DownloadStatus.FAILED)
                return False
            
            self.set_status(DownloadStatus.COMPLETED)
            return success_count > 0
            
        except Exception as e:
            logger.error(f"下载基金详细信息过程中发生异常：{str(e)}")
            return False
    

    
    def _parse_date(self, date_str: Optional[str]) -> Optional[date]:
        """
        解析日期字符串
        """
        if not date_str:
            return None
        
        try:
            # 尝试多种日期格式
            for fmt in ['%Y-%m-%d', '%Y/%m/%d', '%Y年%m月%d日']:
                try:
                    return datetime.strptime(date_str, fmt).date()
                except ValueError:
                    continue
            
            logger.warning(f"无法解析日期格式：{date_str}")
            return None
            
        except Exception as e:
            logger.error(f"解析日期失败：{date_str}, 错误：{str(e)}")
            return None
    
    def download_fund_details_continue(self, funds_queryset) -> bool:
        """
        继续下载基金详细信息
        从指定的基金查询集开始下载
        
        Args:
            funds_queryset: 需要下载的基金查询集
        """
        try:
            # 设置状态和任务信息
            self.set_status(DownloadStatus.DOWNLOADING)
            self._current_task = "基金详细信息继续下载"
            self.reset_stop_flag()
            
            logger.info("开始继续下载基金详细信息...")
            
            # 获取基金代码列表
            fund_codes = list(funds_queryset.values_list('fund_code', flat=True))
            total_count = len(fund_codes)
            
            if total_count == 0:
                logger.info("没有需要下载的基金")
                self.set_status(DownloadStatus.COMPLETED)
                return True
            
            logger.info(f"需要继续下载 {total_count} 只基金的详细信息")
            
            success_count = 0
            failed_count = 0
            
            for i, fund_code in enumerate(fund_codes, 1):
                try:
                    # 检查是否被停止
                    if self.is_stopped():
                        logger.info("下载已被停止")
                        self.set_status(DownloadStatus.STOPPED)
                        return False
                    
                    # 更新进度
                    self._update_progress(
                        f"正在下载第 {i}/{total_count} 只基金详细信息",
                        current_index=i,
                        total_count=total_count,
                        success_count=success_count,
                        failed_count=failed_count
                    )
                    
                    logger.info(f"正在下载第 {i}/{total_count} 只基金：{fund_code}")
                    
                    # 调用akshare接口获取基金详细信息，添加重试机制
                    fund_detail_df = None
                    max_retries = 3
                    
                    for retry in range(max_retries):
                        try:
                            # 检查是否被停止
                            if self.is_stopped():
                                logger.info("下载已被停止")
                                self.set_status(DownloadStatus.STOPPED)
                                return False
                            
                            fund_detail_df = ak.fund_individual_basic_info_xq(symbol=fund_code)
                            break
                        except Exception as retry_error:
                            logger.warning(f"基金 {fund_code} 第 {retry + 1} 次尝试失败：{str(retry_error)}")
                            if retry < max_retries - 1:
                                time.sleep(1)  # 等待1秒后重试
                            else:
                                raise retry_error
                    
                    if fund_detail_df is None or fund_detail_df.empty:
                        logger.warning(f"基金 {fund_code} 详细信息为空")
                        failed_count += 1
                        continue
                    
                    # 解析数据
                    detail_dict = {}
                    for _, row in fund_detail_df.iterrows():
                        detail_dict[row['item']] = row['value']
                    
                    # 创建详细信息记录
                    fund_detail = FundInfoDetail(
                        fund_code=fund_code,
                        fund_name=detail_dict.get('基金名称', ''),
                        fund_full_name=detail_dict.get('基金全称', ''),
                        establishment_date=self._parse_date(detail_dict.get('成立时间', '')),
                        latest_scale=detail_dict.get('最新规模', ''),
                        fund_company=detail_dict.get('基金公司', ''),
                        fund_manager=detail_dict.get('基金经理', ''),
                        custodian_bank=detail_dict.get('托管银行', ''),
                        fund_type=detail_dict.get('基金类型', ''),
                        rating_agency=detail_dict.get('评级机构', ''),
                        fund_rating=detail_dict.get('基金评级', ''),
                        investment_strategy=detail_dict.get('投资策略', ''),
                        investment_objective=detail_dict.get('投资目标', ''),
                        performance_benchmark=detail_dict.get('业绩比较基准', ''),
                        data_source='akshare'
                    )
                    
                    # 保存到数据库
                    fund_detail.save()
                    success_count += 1
                    
                    logger.info(f"基金 {fund_code} 详细信息保存成功")
                    
                    # 避免请求过于频繁
                    time.sleep(0.1)
                    
                except Exception as e:
                    logger.error(f"下载基金 {fund_code} 详细信息失败：{str(e)}")
                    failed_count += 1
                    continue
            
            logger.info(f"基金详细信息继续下载完成，成功：{success_count}，失败：{failed_count}")
            
            # 更新最终进度
            self._update_progress(
                "基金详细信息继续下载完成",
                current_index=total_count,
                total_count=total_count,
                success_count=success_count,
                failed_count=failed_count
            )
            
            if success_count == 0:
                self.set_status(DownloadStatus.FAILED)
                return False
            
            self.set_status(DownloadStatus.COMPLETED)
            return success_count > 0
            
        except Exception as e:
            logger.error(f"继续下载基金详细信息失败：{str(e)}")
            self.set_status(DownloadStatus.FAILED)
            return False
    
    def get_download_progress(self) -> Dict[str, Any]:
        """
        获取下载进度信息
        """
        try:
            # 基金清单进度
            fund_list_count = FundInfo.objects.count()
            
            # 基金详细信息进度
            fund_detail_count = FundInfoDetail.objects.count()
            
            return {
                'fund_list': {
                    'total': fund_list_count,
                    'downloaded': fund_list_count,
                    'progress': 100 if fund_list_count > 0 else 0
                },
                'fund_detail': {
                    'total': fund_list_count,
                    'downloaded': fund_detail_count,
                    'progress': round((fund_detail_count / fund_list_count * 100), 2) if fund_list_count > 0 else 0
                }
            }
            
        except Exception as e:
            logger.error(f"获取下载进度失败：{str(e)}")
            return {}

def main():
    """主函数"""
    downloader = FundDataDownloader()
    
    print("基金数据下载工具")
    print("1. 下载基金清单")
    print("2. 下载基金详细信息")
    print("3. 下载所有数据")
    print("4. 查看下载进度")
    
    choice = input("请选择操作（1-4）：")
    
    if choice == '1':
        success = downloader.download_fund_list()
        print(f"基金清单下载{'成功' if success else '失败'}")
    
    elif choice == '2':
        limit_input = input("请输入下载数量限制（回车表示全部下载）：")
        limit = int(limit_input) if limit_input.strip() else None
        success = downloader.download_fund_details(limit=limit)
        print(f"基金详细信息下载{'成功' if success else '失败'}")
    
    elif choice == '3':
        print("开始下载所有基金数据...")
        
        # 1. 下载基金清单
        print("\n步骤1：下载基金清单")
        if downloader.download_fund_list():
            print("基金清单下载成功")
            
            # 2. 下载基金详细信息（限制前10只用于测试）
            print("\n步骤2：下载基金详细信息")
            if downloader.download_fund_details(limit=10):
                print("基金详细信息下载成功")
            else:
                print("基金详细信息下载失败")
        else:
            print("基金清单下载失败")
    
    elif choice == '4':
        progress = downloader.get_download_progress()
        if progress:
            print("\n下载进度：")
            print(f"基金清单：共计{progress['fund_list']['total']}个，已下载{progress['fund_list']['downloaded']}个，进度{progress['fund_list']['progress']}%")
            print(f"基金详细信息：共计{progress['fund_detail']['total']}个，已下载{progress['fund_detail']['downloaded']}个，进度{progress['fund_detail']['progress']}%")
        else:
            print("获取下载进度失败")
    
    else:
        print("无效选择")

if __name__ == '__main__':
    main()